#include <iostream>
#include <fstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <vector>
#include <functional>

using namespace std;

class VirtualMachine {
public:
    VirtualMachine(const string& filename) {
        // 从文件中读取指令
        ifstream file(filename);
        if (file.is_open()) {
            string line;
            while (getline(file, line)) {
                instructions.push_back(line);
            }
            file.close();
        } else {
            cout << "Unable to open file" << endl;
            exit(1);
        }

        // 初始化指令映射表
        instructionMap.emplace("halt", [this]() { this->halt(); });
        instructionMap.emplace("add", [this]() { this->add(); });
        instructionMap.emplace("sub", [this]() { this->sub(); });
        instructionMap.emplace("mul", [this]() { this->mul(); });
        instructionMap.emplace("div", [this]() { this->div(); });

        instructionMapTwo.emplace("read", [this](string & tokens) { this->read(tokens); });
        instructionMapTwo.emplace("write", [this](string& tokens) { this->write(tokens); });
        instructionMapTwo.emplace("push", [this](string& tokens) { this->push(tokens); });
        instructionMapTwo.emplace("pop", [this](string& tokens) { this->pop(tokens); });
        instructionMapTwo.emplace("label", [this](string& tokens) { this->label(tokens); });
        instructionMapTwo.emplace("jmp", [this](string& tokens) { this->jmp(tokens); });
        instructionMapTwo.emplace("jnz", [this](string &tokens) { this->jnz(tokens); });
        instructionMapTwo.emplace("jeg", [this](string &tokens) { this->jeg(tokens); });
        instructionMapTwo.emplace("jel", [this](string &tokens) { this->jel(tokens); });
    }

    void run() {
        // 初始化label
        for(int i=0;i<instructions.size();i++){
            if(instructions[i].find("label")!=string::npos){
                labels[instructions[i].substr(6)]=i;
            }
        }
        // 运行虚拟机
        while (pc < instructions.size()) {
            
            execute(instructions[pc]);
            pc++;
        }
    }

private:
    stack<int> stack_;
    unordered_map<string, function<void()>> instructionMap;
    unordered_map<string,function<void(string &)>> instructionMapTwo;
    unordered_map<string, int> labels;
    unordered_map<string,int> variables;
    vector<string> instructions;
    int pc = 0; // 程序计数器

    void halt() {
        // 停止运行
        exit(0);
    }

    void read(string & tokens) {
        // 从输入流中读取一个值并存入寄存器
        cout << "Enter a value for variable " << tokens << ": ";
        int value;
        cin >> value;
        variables[tokens] = value;
    }

    void write(string &tokens) {
        // 将寄存器的值输出到输出流
        cout <<tokens << " = " << variables[tokens] << endl;
    }

    void push(string& tokens) {
        // 将寄存器或常数压入栈顶
        if (tokens.front() == '$') {
            int reg = stoi(tokens.substr(1));
            stack_.push(reg);
        } else if(variables.count(tokens)==1){
            stack_.push(variables[tokens]);
        }else{
            cerr<<"Undefined variable: "<<tokens<<endl;
        }
    }

    void pop(string& tokens) {
        // 将栈顶的值弹出并存入寄存器
        int reg = stack_.top();
        variables[tokens] = reg;
        stack_.pop();
    }

    void label(string& tokens) {
        // 啥也不用作，前面记录过了
  
    }

    void jmp(string &tokens) {
        // 跳转到标签
        if(labels.count(tokens)==0){
            cerr << "Undefined label: " << tokens << endl;
            exit(1);
        }
        pc = labels[tokens];
    }

    void jnz(string &tokens) {
        // 如果栈顶的值为0，则跳转到标签
        if (stack_.top() != 0) {
            if(labels.count(tokens)==0){
                cerr << "Undefined label: " << tokens << endl;
                exit(1);
            }
            pc = labels[tokens];
        }
        stack_.pop();
    }

    void jeg(string &tokens){
        // 如果栈顶的值大于等于0，则跳转到标签
        if (stack_.top() >= 0) {
            if(labels.count(tokens)==0){
                cerr << "Undefined label: " << tokens << endl;
            }
            pc=labels[tokens];
        }
        stack_.pop();
    }

    void jel(string &tokens){
        // 如果栈顶的值小于等于0，则跳转到标签
        if (stack_.top() <= 0) {
            if(labels.count(tokens)==0){
                cerr << "Undefined label: " << tokens << endl;
            }
            pc=labels[tokens];
        }
        stack_.pop();
    }

    void add() {
        // 栈顶两个元素相加
        int a = stack_.top();
        stack_.pop();
        int b = stack_.top();
        stack_.pop();
        stack_.push(b + a);
    }

    void sub() {
        // 栈顶两个元素相减
        int a = stack_.top();
        stack_.pop();
        int b = stack_.top();
        stack_.pop();
        stack_.push(b - a);
    }

    void mul() {
        // 栈顶两个元素相乘
        int a = stack_.top();
        stack_.pop();
        int b = stack_.top();
        stack_.pop();
        stack_.push(a * b);
    }

    void div() {
        // 栈顶两个元素相除
        int a = stack_.top();
        stack_.pop();
        int b = stack_.top();
        stack_.pop();
        stack_.push(b / a);
    }

    void execute(string& tokens) {
        if(tokens.front()=='#'){
            return ; // 注释
        }
        // 执行指令
        if (instructionMap.count(tokens)) {
            instructionMap[tokens]();
        } else {
            auto posi= tokens.find(' ');
            if(posi != string::npos){
                string instr= tokens.substr(0,posi);
                string arg= tokens.substr(posi+1);
                if(instructionMapTwo.count(instr)){
                    instructionMapTwo[instr](arg);
                }else{

                    cout << "Invalid instruction: " << tokens << endl;
                }
            }else{
                cout << "Invalid instruction: " << tokens << endl;
            }
        }
    }
};

int main(int argc, char* argv[]) {
    if (argc != 2) {
        cout << "Usage: " << argv[0] << " <input_file>" << endl;
        return 1;
    }

    VirtualMachine vm(argv[1]);
    vm.run();
    return 0;
}